### File Purpose: Produce the maps in Figure 3 and Figure A3 ###

# Get the maps from the tigris package
counties_map <- counties(cb = TRUE) %>%
  st_as_sf() %>%
  st_transform("+init=epsg:2163") %>%
  mutate(fips = as.numeric(paste0(STATEFP, COUNTYFP)))

states_map <- states(cb = TRUE) %>%
  st_as_sf() %>%
  st_transform("+init=epsg:2163")

# Get the COVID and population data
covid_dat_in <-
  read_csv("../data/derived/cases/usa_cases_deaths_jh_nyt_unique.csv")

pop_dat_in <- 
  read_csv("../data/derived/covariates/county_covariates.csv")

#################################################
## 1. Map of Exposure on March 15 -- Figure A3 ##
#################################################

covid_dat_march15 <-
  covid_dat_in %>% 
  filter(date == "2020-03-15")

### PULL OUT STAT FOR TEXT ###
covid_dat_march15 %>% filter(fips >= 1000, fips < 72000) %>% .$cases_jh %>% mean
covid_dat_march15 %>% filter(fips >= 1000, fips < 72000) %>% .$cases_jh %>% sd

# Merge with shape files
dat_map_march15 <- 
  right_join(covid_dat_march15,
             counties_map,
             by=c("fips")) %>% 
  st_as_sf

# Create clean buckets for these levels
dat_map_march15 <- dat_map_march15 %>%
  mutate(cases_bkt = case_when(
    cases_jh < 1 ~ "0",
    cases_jh < 10 ~ "1-9",
    cases_jh < 25 ~ "10-24",
    cases_jh < 100 ~ "25-100",
    cases_jh >= 100 ~ "100+")) %>%
  mutate(cases_bkt = factor(cases_bkt, levels=c("0", "1-9", "10-24", "25-100", "100+")))

# Plot
ggplot(filter(dat_map_march15, !is.na(cases_bkt))) +
  geom_sf(aes(fill = cases_bkt), colour="#ADADAD", lwd=0) +
  geom_sf(data=states_map, fill="transparent", colour="#A1A1A1", size=0.2) +
  theme_void() +
  labs(fill = "Cases") +
  scale_fill_brewer(palette = "YlOrRd", drop=FALSE) +
  theme(legend.text  = element_text(size = 8),
        legend.key.size = unit(0.8, "lines"),
        legend.position = "bottom", legend.box = "horizontal") +
  guides(fill = guide_legend(nrow = 1, title.hjust = 0.5)) +
  coord_sf(xlim = c(-2200000, 2700000), ylim = c(-2200000, 850000), expand = FALSE) 


ggsave(str_interp("../plots/march15_covid_cases.jpg"),
       width = 6, height = 4,  units = "in", last_plot())



####################################################
## 2. Map of Change in Cases by Month -- Figure 3 ##
####################################################

dates <- tibble(
  key = c("Feb", "March", "April", "May", "June", "July"),
  date = as.Date(c("2020-02-28", "2020-03-27", "2020-04-24", "2020-05-29", "2020-06-26", "2020-07-24")))

covid_dat <- dates %>%
  left_join(covid_dat_in, by="date") %>%
  inner_join(pop_dat_in, by=c("fips"="county")) %>%
  mutate(
    cases = cases_jh,
    cases_per_100k = cases_jh/(tot_pop_2018/100000)) %>%
  mutate(log_cases_per_100k = log(cases_per_100k + 1)) %>%
  group_by(fips) %>%
  # Ensure every county has all 6 observations
  mutate(n = n()) %>%
  filter(n == 6) %>%
  # Get the change measure
  mutate(change_log_cases_per_100k = log_cases_per_100k - lag(log_cases_per_100k)) %>%
  ungroup %>%
  # Look at percentiles of map
  group_by(key) %>%
  mutate(ptile_change_log_cases_per_100k = ntile(change_log_cases_per_100k, 100)) %>%
  ungroup()

dat_map <- 
  right_join(covid_dat,
             counties_map,
             by=c("fips")) %>% 
  st_as_sf


graphs_info <- 
  tibble(month = c("March",
                   "April",
                   "May",
                   "June",
                   "July"),
         label_nrml = c("%-tile March - Feb, log(Cases per 100k)",
                   "%-tile  April - March, log(Cases per 100k)",
                   "%-tile May - April, log(Cases per 100k)",
                   "%-tile June - May, log(Cases per 100k)",
                   "%-tile July - June, log(Cases per 100k)")
  )

## Loop through the months and produce the final maps
for(i in 1:nrow(graphs_info)){
  
  month <- graphs_info$month[i]
  label <- graphs_info$label_nrml[i]
  
  ggplot(filter(dat_map, key==month)) +
    geom_sf(aes(fill = ptile_change_log_cases_per_100k), colour="#cccccc", lwd=0) +
    labs(fill = label) +
    theme_void() +
    scale_fill_binned(type = "viridis",
                        breaks = c(0, 50, 75, 90, 95),
                        limits = c(0, 100),
                        guide = guide_coloursteps(even.steps = FALSE,
                                                  show.limits = TRUE)) +
    theme(legend.text  = element_text(size = 8),
          legend.key.size = unit(0.8, "lines"),
          legend.position = "bottom", legend.box = "horizontal") +
    guides(fill = guide_legend(nrow = 1, title.hjust = 0.5)) +
    coord_sf(xlim = c(-2200000, 2700000), ylim = c(-2200000, 850000), expand = FALSE)
  
  ggsave(str_interp("../plots/${month}_change_log_cases_per_100k.jpg"),
         width = 5, height = 3,  units = "in", last_plot())
  
}